from pathlib import Path
import random


def make_data_set():
    path_train = Path(r'F:\datasets\breast\labelimg\voc\images')  # --图片数据集目录
    path_dataset = Path(r'D:\mypro\yolov8\dataset\breast')  # --存放数据集的划分文件
    f_train = open(path_dataset / 'train.txt', 'w')
    f_valid = open(path_dataset / 'val.txt', 'w')
    f_test = open(path_dataset / 'test.txt', 'w')
    f_train_val = open(path_dataset / 'trainval.txt', 'w')
    for path in path_train.glob('*.*'):
        i = random.random()
        if i < 0.1:
            f_test.write(str(path) + '\n')
        elif i < 0.2:
            f_valid.write(str(path) + '\n')
        else:
            f_train.write(str(path) + '\n')
        f_train_val.write(str(path) + '\n')
    f_train.close()
    f_valid.close()
    f_test.close()
    f_train_val.close()


if __name__ == "__main__":
    make_data_set()
